import pandas as pd
import torch
from datasets import Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, Trainer
from modelscope import snapshot_download
from peft import LoraConfig, TaskType, get_peft_model, PeftModel
import os


# 需要微调的基座模型
model_id = 'DeepSeek-R1-Distill-Llama-8B'

# 检查CUDA是否可用，然后检查MPS是否可用，最后回退到CPU
device = torch.device("cuda:0")

base_dir = '/app'
models_dir = f'{base_dir}/models'
dataset_file = f'{base_dir}/datasets/huanhuan.json'
model_path = f"{models_dir}/{model_id.replace('.', '___')}"
checkpoint_dir = f"./checkpoint/{model_id}"
lora_dir = f"./lora/{model_id}"


# 打印路径（用于验证）
print(f"Models Directory: {models_dir}")
print(f"Dataset File: {dataset_file}")
print(f"Model Path: {model_path}")
print(f"Checkpoint Directory: {checkpoint_dir}")
print(f"LoRA Directory: {lora_dir}")
torch_dtype = torch.half

lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1
)

# 新增的推理函数，只接收一个输入文本参数
def generate_output(input_text):
    # 加载 tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token

    # 加载基础模型
    base_model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch_dtype)
    base_model.to(device)

    # 加载LoRA权重
    lora_model = PeftModel.from_pretrained(base_model, lora_dir)
    lora_model.eval()

    # 构建输入文本，省略 Complex_CoT 部分
    instruction = "请具体分析"
    gen_text = f"<|start_header_id|>user<|end_header_id|>\n\n{instruction}: {input_text}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
 
    # 构建输入文本，使用简单提示替代 Complex_CoT
    #gen_text = f"<|start_header_id|>user<|end_header_id|>\n\n{instruction}: {input_text}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n请逐步分析这个病例。"
    
    # 动态生成简单思维链
    #simple_cot = "首先考虑症状、检查结果，然后分析可能的疾病类型。"
    #gen_text = f"<|start_header_id|>user<|end_header_id|>\n\n{instruction}: {input_text}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n{simple_cot}"


    input_ids = tokenizer(gen_text, return_tensors="pt").input_ids.to(device)

    # 生成输出
    with torch.no_grad():
        outputs = lora_model.generate(input_ids=input_ids, max_new_tokens=300, do_sample=True, top_p=0.9, temperature=0.8)

    # 解码输出
    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return output_text



def test():
    #train()
    res = generate_output("你是谁？")
    print(res)

# 新增的推理函数，只接收一个输入文本参数
def predict(input_text, tokenizer, lora_model):

# 构建输入文本模板
    instruction = "请具体分析"
    gen_text = f"<|start_header_id|>user<|end_header_id|>\n\n{instruction}: {input_text}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
    
    # 对输入文本进行编码，同时生成 attention_mask
    encoding = tokenizer(gen_text, return_tensors="pt", padding=True, truncation=True)
    
    # 将输入移动到相同的设备（CPU或GPU）
    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)

    # 生成输出
    with torch.no_grad():
        outputs = lora_model.generate(input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, top_p=0.9, temperature=0.8, pad_token_id=tokenizer.eos_token_id)

    # 解码输出
    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return output_text


def test_run():
    # 初始化模型、tokenizer等资源（假设这些资源初始化只需要一次）
    # 注意：这部分应包含所有一次性初始化的代码，如加载model和tokenizer等
    
    print("输入您的问题（或使用 Ctrl+D 来退出）:")
    try:
        while True:
            # 获取用户输入
            input_text = input("您: ")
            
            # 如果输入为空，则跳过
            if not input_text.strip():
                continue
            
            # 调用函数生成回答
            res = generate_output(input_text)
            
            # 输出回答
            print(f"{res}")
            
    except EOFError:
        # 当用户使用 Ctrl+D (Unix/Linux) 或 Ctrl+Z (Windows) 退出时，捕获EOFError异常
        print("\n对话结束")

if __name__ == '__main__':
    # 初始化模型、tokenizer等资源（假设这些资源初始化只需要一次）
    # 注意：这部分应包含所有一次性初始化的代码，如加载model和tokenizer等
     # 加载 tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token

    # 加载基础模型
    base_model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch_dtype)
    base_model.to(device)

    # 加载LoRA权重
    lora_model = PeftModel.from_pretrained(base_model, lora_dir)
    lora_model.eval()   
    print("输入您的问题（或使用 Ctrl+D 来退出）:")
    try:
        while True:
            # 获取用户输入
            input_text = input("您: ")
            
            # 如果输入为空，则跳过
            if not input_text.strip():
                continue
            
            # 调用函数生成回答
            res = predict(input_text, tokenizer, lora_model)
            
            # 输出回答
            print(f"{res}")
            
    except EOFError:
        # 当用户使用 Ctrl+D (Unix/Linux) 或 Ctrl+Z (Windows) 退出时，捕获EOFError异常
        print("\n对话结束")
